'''
Created on Oct 5, 2010

@author: oabalbin
'''
import sys
import operator
import numpy as np
import scipy as sp
import random
from scipy import stats
from collections import defaultdict, deque


def copa_analysis(expmatrix, median_shift, ids_groupa, ids_groupb, percentiles, sdv_factor=0.25, mmean=True):
    '''
    expmatrix expression matrix
    expression_cutoff: expression normalization factor
    ids_groupa: tipically normal samples, ids_groupb: tumor samples
    returns mean_gene_copa_score_list (list of outlier genes), copa_score_array (array with all copa scores), 
    gene_outliersample_dict (samples for which gene i is an outlier
    '''
    
    sample_indicator = np.array(range(expmatrix.shape[1]))
    gene_outliersample_dict=defaultdict(list)
    copa_score_array = np.empty(expmatrix.shape[0])
    outliers_list=[]
    matrix_outliers_bysample = np.zeros((expmatrix.shape[0],expmatrix.shape[1]))
    
    for i in range(expmatrix.shape[0]):
        # Get row statistics
        sdt_normals = np.std(expmatrix[i,ids_groupa]) 
        sdt_tumors = np.std(expmatrix[i,ids_groupb])
        gene_median = np.median(expmatrix[i,:])
        # Get the gene mad
        gene_mad = np.median(np.abs(expmatrix[i,:] - gene_median)) + median_shift
        copa_normalized_row = (expmatrix[i,:] - gene_median)/gene_mad
        # copa scores
        copa_scores = [np.array(sp.stats.scoreatpercentile(copa_normalized_row,per)) for per in percentiles]
        # get the mean copa score
        if mmean: 
            mean_copa_score = np.mean(copa_scores)
        else: 
            mean_copa_score = np.median(copa_scores)
            
        # Get outlier_samples 
        outlier_samples = sample_indicator[copa_normalized_row > mean_copa_score]
        false_positives = set(outlier_samples).intersection(set(ids_groupa))
        gene_outliersample_dict[i].append(outlier_samples)
        # copa score to report: copa score*number of outlier samples
        wmean_copa_score = mean_copa_score*len(outlier_samples)
        copa_score_array[i] = wmean_copa_score
        '''
        if i == 62322:
            print outlier_samples
            print mean_copa_score, len(outlier_samples), copa_scores
            sys.exit(0)
        '''
        # Filter outliers according to the level of variance in the normals
        # if (len(outlier_samples) <=2) or (sdt_normals > sdt_tumors*sdv_factor) or (len(false_positives) > 2):
        if sdt_normals > sdt_tumors*sdv_factor:
            continue
        # Make a list of outliers 
        outliers_list.append( type('geneCopa', (object,), 
                                   dict(indicator=i,
                                        mean_copa_score=wmean_copa_score, 
                                        outlier_samples=outlier_samples,
                                        median_exp = gene_median,
                                        max_exp = np.max(expmatrix[i,:]))) )
        
        # Fill the matrix of outlier by sample
        matrix_outliers_bysample[i,outlier_samples] = wmean_copa_score
        
    #sorted list of outliers
    outliers_list = sorted(outliers_list, key=operator.attrgetter('mean_copa_score'), reverse=True)
    
    return outliers_list, copa_score_array, gene_outliersample_dict, matrix_outliers_bysample 


def calculate_score_pvalue(copa_scores, numberOfpermutations):
    '''
    It returns the pvalue for each row
    copa_scores,permutations_number
    '''
    
    random_copa_scores = np.array(random.sample(copa_scores,numberOfpermutations))
    score_pvals = np.empty(len(copa_scores))
    for i in range(len(copa_scores)):
        score_pvals[i]= len(random_copa_scores[random_copa_scores >= copa_scores[i]])/float(numberOfpermutations)
            
    return score_pvals
    

def print_list_of_outliers(list_outliers_report, outfile_outliers_list, random_copa_scores, gene_names, gene_locations, \
                           knownHugoNames=[], cohort_sizes=[0,0,0]):
    """
    Prints the list of outliers
    """       
    header = [ "rank", "gene.indicator", "gene.name", "gene.location", "gene.mean_copa_score", "Nsamples", "max_exp","median_exp", "gene.pvalue"] 
    outfile_outliers_list.write("cohort,tumor,benign"+'\t'+",".join(map(str,cohort_sizes)).replace(',','\t')+'\n')
    outfile_outliers_list.write(",".join(map(str,header)).replace(',','\t')+'\n')
    
    
    for rank, tho in enumerate(list_outliers_report):
        if knownHugoNames:
            genename=gene_names[tho.indicator]+'_'+knownHugoNames[gene_names[tho.indicator]][0][0]
        else:
            genename=gene_names[tho.indicator]
            
        pvalue=random_copa_scores[tho.indicator]
        line = [rank, tho.indicator, genename, gene_locations[tho.indicator],\
                tho.mean_copa_score, len(tho.outlier_samples), tho.max_exp, tho.median_exp, pvalue]        
        #print ",".join(map(str,line)).replace(',','\t')
        outfile_outliers_list.write(",".join(map(str,line)).replace(',','\t')+'\n')
        

def print_matrix_of_outliers(samples_dict, gene_dict, outliers2report, matrix_outliers_bysample, outfile_outliers_matrix):
    """
    Prints the matrix of outliers
    """       
    list_of_samples = get_list_of_samples(samples_dict)
    outfile_outliers_matrix.write("geneName"+'\t'+",".join(list_of_samples).replace(',','\t')+'\n')
    
    for tho in outliers2report:
        eline = list(matrix_outliers_bysample[tho.indicator,:])
        outfile_outliers_matrix.write(gene_dict[tho.indicator]+'\t'+",".join(map(str,eline)).replace(',','\t')+'\n') 


def get_list_of_samples(sampledict):
    print_samples=['z' for i in range(len(sampledict))]
    for j, sp in sampledict.iteritems():
        print_samples[j] = sp
    return list(print_samples)



# Calc copa on matrix all at once. 10-8-10: To be deprecated because all analysis can be done at the row level
def copa_norm_across_samples(expmatrix, median_shift):
    """
    Copa analisis accros the different samples
    """
    
    gene_median = np.median(expmatrix,axis=1)
    median_col = np.reshape(gene_median,(expmatrix.shape[0],-1))
    gene_median_matrix = np.tile(median_col,expmatrix.shape[1])
    # Get the gene mad
    gene_mad = np.median(np.abs(expmatrix - gene_median_matrix),axis=1) + median_shift
    ############# Important ################
    # when gene_mad = 0.0 set it to den_mad = 1.0
    #gene_mad[gene_mad == 0.0] = 1.0
     
    # make a matrix with the mad
    gene_mad_col = np.reshape(gene_mad,(expmatrix.shape[0],-1))
    gene_mad_matrix = np.tile(gene_mad_col,expmatrix.shape[1])
    # Get the copa index for each gene in each cell line
    copa_normalized_matrix = np.true_divide((expmatrix - gene_median_matrix), gene_mad_matrix)
    
    return copa_normalized_matrix

## Permutation routine. 10-8-10: To be deprecated because is to slow
def permute_matrix_indexes(nlabels, nperm):
    """
    It outputs a matrix of permutations, 
    """
    return np.random.random_integers(0,nlabels-1, size=(nlabels, nperm))


def copa_permute_matrix_rows(expmatrix,id_groupa, id_groupb, npermutations, percentils_list, expression_cutoff, sdv_factor, method):
    
    randlabels = permute_matrix_indexes(expmatrix.shape[0], npermutations)
    random_copa_scores = np.zeros((randlabels.shape[1],randlabels.shape[0]))
    
    for i in range(randlabels.shape[1]):
        print "Permutations", i
        permuted_expmatrix = expmatrix[randlabels[:,i],:]
        # copa normalization for the permute matrix
        permuted_copa_matrix = copa_norm_across_samples(permuted_expmatrix,expression_cutoff)
        gene_outliersample_dict, copa_rank_dict, copa_score_dict = copa_at_different_percentiles(permuted_copa_matrix,percentils_list)
        # calculate the copa score for each row.
        random_copa_scores[i,:] = summarize_copa_results_permuted(permuted_copa_matrix, id_groupa, id_groupb, gene_outliersample_dict, copa_score_dict, sdv_factor, method)
        
    return np.transpose(random_copa_scores)
        
    
def summarize_copa_results_permuted(expmatrix, id_groupa, id_groupb, gene_outliersample_dict, copa_score_dict, sdv_factor, method='mean'):
    """
    Returns a rank list of final outliers.
    """
    indicator=np.array(range(len(copa_score_dict[0])))
    gene_copascore_values = np.empty(len(copa_score_dict))    
    basal_level = 0.0
    
    for i in  range(len(copa_score_dict)):
        
        if method == 'median':
            mean_copa_score = np.median(copa_score_dict[i])
        elif method == 'mean':
            mean_copa_score = np.mean(copa_score_dict[i])
        mean_copa_ind = indicator[copa_score_dict[i]>=mean_copa_score][0]
        outlier_samples = gene_outliersample_dict[i][mean_copa_ind]
        
        # Filter outliers according to the level of variance in the normals
        sdt_normals = np.std(expmatrix[i,id_groupa]) 
        sdt_tumors = np.std(expmatrix[i,id_groupb])

        if sdt_normals > sdt_tumors*sdv_factor:
            gene_copascore_values[i] = basal_level
        else:
            mean_copa_score = mean_copa_score*len(outlier_samples)
            # Make a temporal array with only the copa score for each gene
            gene_copascore_values[i] = mean_copa_score
            
    return gene_copascore_values

    
def calculate_score_pvalue_mat(random_copa_matrix, copa_scores, numberOfpermutations):
    '''
    It returns the pvalue for each row
    random_copa_matrix, copa_scores
    permutations_number
    '''
    score_pvals = np.empty(len(copa_scores))
    for i in range(len(copa_scores)):
        score_pvals[i]= len(random_copa_matrix[i,:] >= copa_scores[i])/float(numberOfpermutations)
    
    return score_pvals
